# 在文件顶部添加导入
import sys
import os
import time
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QLineEdit, QPushButton, QFileDialog, QTextEdit,
                             QComboBox, QTabWidget, QMessageBox, QGroupBox, QRadioButton, QListWidget)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from custom_crypto_system import crypto_system
from p2p_encrypted_communication import p2p_communication

# 修改主窗口类，添加点对点通信选项卡
class CryptoGUI(QMainWindow):
    # 添加信号用于在主线程中更新UI
    message_received = pyqtSignal(dict)
    
    def __init__(self):
        super().__init__()
        self.setWindowTitle("自定义密码加密系统")
        self.setGeometry(100, 100, 800, 600)
        
        # 连接信号
        self.message_received.connect(self.update_message_display)
        
        # 创建主窗口部件
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        
        # 创建主布局
        self.main_layout = QVBoxLayout()
        self.central_widget.setLayout(self.main_layout)
        
        # 创建选项卡
        self.tab_widget = QTabWidget()
        self.main_layout.addWidget(self.tab_widget)
        
        # 创建各个选项卡
        # self.create_password_crypto_tab()  # 密码加密选项卡
        # self.create_rsa_tab()              # RSA加密选项卡
        # self.create_keygen_tab()           # 密钥生成选项卡
        self.create_p2p_tab()              # 新增点对点通信选项卡

    # 在现有代码后添加以下方法
    def create_p2p_tab(self):
        """创建点对点加密通信选项卡"""
        p2p_widget = QWidget()
        p2p_layout = QVBoxLayout()
        p2p_widget.setLayout(p2p_layout)
        
        # 服务器控制组
        server_group = QGroupBox("服务器控制")
        server_layout = QVBoxLayout()
        server_group.setLayout(server_layout)
        
        # 本地IP和端口
        ip_port_layout = QHBoxLayout()
        self.local_ip_label = QLabel("本地IP: 检测中...")
        port_layout = QHBoxLayout()
        port_layout.addWidget(QLabel("端口:"))
        self.server_port = QLineEdit("8080")
        port_layout.addWidget(self.server_port)
        port_layout.addStretch()
        
        ip_port_layout.addWidget(self.local_ip_label)
        ip_port_layout.addLayout(port_layout)
        server_layout.addLayout(ip_port_layout)
        
        # 服务器按钮
        server_btns_layout = QHBoxLayout()
        self.start_server_btn = QPushButton("启动服务器")
        self.start_server_btn.clicked.connect(self.start_p2p_server)
        self.stop_server_btn = QPushButton("停止服务器")
        self.stop_server_btn.clicked.connect(self.stop_p2p_server)
        self.stop_server_btn.setEnabled(False)
        
        server_btns_layout.addWidget(self.start_server_btn)
        server_btns_layout.addWidget(self.stop_server_btn)
        server_layout.addLayout(server_btns_layout)
        
        p2p_layout.addWidget(server_group)
        
        # 对等节点连接组
        peer_group = QGroupBox("连接对等节点")
        peer_layout = QVBoxLayout()
        peer_group.setLayout(peer_layout)
        
        # 对等节点IP和端口
        peer_connect_layout = QHBoxLayout()
        peer_connect_layout.addWidget(QLabel("对等节点IP:"))
        self.peer_ip = QLineEdit()
        peer_connect_layout.addWidget(self.peer_ip)
        peer_connect_layout.addWidget(QLabel("端口:"))
        self.peer_port = QLineEdit("8080")
        peer_connect_layout.addWidget(self.peer_port)
        connect_btn = QPushButton("连接")
        connect_btn.clicked.connect(self.connect_to_peer)
        peer_connect_layout.addWidget(connect_btn)
        # 断开连接按钮
        disconnect_btn = QPushButton("断开连接")
        disconnect_btn.clicked.connect(self.disconnect_peer)
        peer_connect_layout.addWidget(disconnect_btn)
        peer_layout.addLayout(peer_connect_layout)
        
        
        # 已连接对等节点列表
        peer_list_layout = QVBoxLayout()
        peer_list_layout.addWidget(QLabel("已连接的对等节点:"))
        self.peer_list = QListWidget()
        peer_list_layout.addWidget(self.peer_list)
        peer_layout.addLayout(peer_list_layout)
        
        p2p_layout.addWidget(peer_group)
        
        # 消息发送和接收组
        message_group = QGroupBox("加密通信")
        message_layout = QVBoxLayout()
        message_group.setLayout(message_layout)
        
        # 消息显示区域
        self.message_display = QTextEdit()
        self.message_display.setReadOnly(True)
        self.message_display.setPlaceholderText("通信消息将显示在这里...")
        message_layout.addWidget(QLabel("消息记录:"))
        message_layout.addWidget(self.message_display)
        
        # 消息发送区域
        message_send_layout = QHBoxLayout()
        self.message_input = QLineEdit()
        self.message_input.setPlaceholderText("输入要发送的消息")
        send_btn = QPushButton("发送")
        send_btn.clicked.connect(self.send_message)
        message_send_layout.addWidget(self.message_input)
        message_send_layout.addWidget(send_btn)
        message_layout.addLayout(message_send_layout)
        
        p2p_layout.addWidget(message_group)
        
        self.tab_widget.addTab(p2p_widget, "点对点加密通信")
        
        # 获取并显示本地IP
        self.update_local_ip()
    
    def update_local_ip(self):
        """更新本地IP显示"""
        try:
            local_ip = p2p_communication.get_local_ip()
            self.local_ip_label.setText(f"本地IP: {local_ip}")
        except Exception as e:
            self.local_ip_label.setText(f"获取IP失败: {str(e)}")
    
    def start_p2p_server(self):
        """启动P2P服务器"""
        try:
            port = int(self.server_port.text())
            p2p_communication.message_callback = self.on_message_received
            p2p_communication.start_server(port)
            
            self.start_server_btn.setEnabled(False)
            self.stop_server_btn.setEnabled(True)
            self.update_message_display({"from": "系统", "message": f"服务器已启动，监听端口 {port}", "timestamp": time.time()})
        except Exception as e:
            QMessageBox.critical(self, "错误", f"启动服务器失败: {str(e)}")
    
    def stop_p2p_server(self):
        """停止P2P服务器"""
        try:
            p2p_communication.stop_server()
            
            self.start_server_btn.setEnabled(True)
            self.stop_server_btn.setEnabled(False)
            self.update_message_display({"from": "系统", "message": "服务器已停止", "timestamp": time.time()})
        except Exception as e:
            QMessageBox.critical(self, "错误", f"停止服务器失败: {str(e)}")
    
    def connect_to_peer(self):
        """连接到对等节点"""
        peer_ip = self.peer_ip.text()
        try:
            peer_port = int(self.peer_port.text())
            if not peer_ip:
                QMessageBox.warning(self, "警告", "请输入对等节点IP")
                return
            
            success = p2p_communication.connect_to_peer(peer_ip, peer_port)
            if success:
                peer_info = f"{peer_ip}:{peer_port}"
                self.peer_list.addItem(peer_info)
                self.update_message_display({"from": "系统", "message": f"成功连接到对等节点: {peer_info}", "timestamp": time.time()})
            else:
                QMessageBox.warning(self, "警告", "连接到对等节点失败")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"连接失败: {str(e)}")
    
    def disconnect_peer(self):
        """断开与对等节点的连接"""
        selected_items = self.peer_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "警告", "请先选择要断开的对等节点")
            return
        
        peer_info = selected_items[0].text()
        p2p_communication.disconnect_peer(peer_info.split(":")[0], int(peer_info.split(":")[1]))
        self.peer_list.takeItem(self.peer_list.row(selected_items[0]))
        self.update_message_display({"from": "系统", "message": f"已断开与 {peer_info} 的连接", "timestamp": time.time()})
    def send_message(self):
        """发送加密消息"""
        message = self.message_input.text()
        if not message:
            QMessageBox.warning(self, "警告", "请输入要发送的消息")
            return
        
        # 获取选中的对等节点
        selected_items = self.peer_list.selectedItems()
        if not selected_items:
            QMessageBox.warning(self, "警告", "请先选择要发送的对等节点")
            return
        
        try:
            peer_info = selected_items[0].text()
            peer_ip, peer_port = peer_info.split(":")
            success = p2p_communication.send_message(peer_ip, int(peer_port), message)
            
            if success:
                self.update_message_display({"from": "我", "message": message, "timestamp": time.time()})
                self.message_input.clear()
            else:
                QMessageBox.warning(self, "警告", "发送消息失败")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"发送消息失败: {str(e)}")
    
    def on_message_received(self, message_data):
        """接收到消息的回调函数"""
        # 发送信号到主线程更新UI
        self.message_received.emit(message_data)
    
    def update_message_display(self, message_data):
        """更新消息显示区域"""
        timestamp = time.strftime('%H:%M:%S', time.localtime(message_data['timestamp']))
        sender = message_data['from']
        message = message_data['message']
        
        # 添加新消息到显示区域
        self.message_display.append(f"[{timestamp}] {sender}: {message}")
        # 滚动到底部
        self.message_display.verticalScrollBar().setValue(self.message_display.verticalScrollBar().maximum())

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = CryptoGUI()
    window.show()
    sys.exit(app.exec_())